package org.ykhl.mq.commons;

/**
 * Created by zkf on 12/29/15.
 */
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

public class AbortableCountDownLatch extends CountDownLatch {
    protected volatile boolean aborted = false;
    private volatile String msg = null;

    public AbortableCountDownLatch(int count) {
        super(count);
    }

    /**
     * Unblocks all threads waiting on this latch and cause them to receive an
     * AbortedException.  If the latch has already counted all the way down,
     * this method does nothing.
     */
    public void abort(String msg) {
        if( getCount()==0 )
            return;

        this.aborted = true;
        if(msg != null) this.msg = msg;
        while(getCount()>0)
            countDown();
    }

    @Override
    public boolean await(long timeout, TimeUnit unit) throws InterruptedException {
        final boolean rtrn = super.await(timeout,unit);
        throwIfAborted();
        return rtrn;
    }

    @Override
    public void await() throws InterruptedException {
        super.await();
        throwIfAborted();
    }

    private void throwIfAborted() throws AbortedException {
        if (aborted){
            if(msg == null)
                throw new AbortedException();
            else
                throw new AbortedException(msg);
        }
    }

    public static class AbortedException extends InterruptedException {
        private static final long serialVersionUID = 5650113237373433118L;

        public AbortedException() {
        }

        public AbortedException(String detailMessage) {
            super(detailMessage);
        }
    }
}
